import time
import BaseHTTPServer
import uuid
import random
import os
import re
import logging
import shutil
import threading
import time

# Simple BITS Upload Server implementation

# Some notes :

# - Does not implement the 413 message regarding fragment size  - couldn't find evidence of others doing so and it made no odds
# - Does not implement the Content-Name header as documentation suggests this is not used
# - Does not allow for subfolders of virtual directories - specify as many virtual directories as you like, but uploads must be in that folder
# - Does not implement Download - just add this to an existing HTTP server
# - Cleans up incomplete sessions after a timeout
# - Does not implement the UPLOAD-REPLY message at this time, just simple upload

# See the __main__ at the bottom of this file for how it is used. It is not intended to run standalone, but more to have parts taken from
# it for implementation into a fuller HTTP server

# Tried on Windows 7 and Ubuntu 13.04

class BitsSessionCreationError(Exception):
    """ Error thrown when an issue is encountered trying to create a temporary session directory for a new upload session """

    def __init__(self, reason):
        self.reason = reason

    def __str__(self):
        return repr(self.reason)

class BitsFragmentSyncError(Exception):
    """ Error thrown when client and server ranges get out of sync, leading to returning an HTTP 416 error """

    def __init__(self, nextExpectedByte):
        self.nextExpectedByte = nextExpectedByte

    def __str__(self):
        return repr(self.nextExpectedByte)

class BitsSessionCleaner(threading.Thread):
    """ Looks through session logs to find any partially completed ones that need deleting """
    """ Run as a separate thread """

    def __init__(self, config):
        super(BitsSessionCleaner, self).__init__()
        self._stop = threading.Event()
        self.config = config

    def run(self):
        while not self._stop.isSet():
            self._stop.wait(2)

            transitDir = self.config["transit"]
            for root, dirs, files in os.walk(transitDir, topdown=False):
                for name in files:
                    if str(name).endswith("bits-data"):                                # Test the last time data was written to this file
                        filepath = os.path.join(root, name)
                        lastModified = os.path.getmtime(filepath)
                        curTime = time.time()

                        timeDiff = curTime - lastModified

                        if timeDiff > self.config["timeout"]:                        # If the file hasn't been written to since for longer
                            try:                                                    # than the timeout, try to delete the session folder
                                shutil.rmtree(root)
                            except OSError:
                                pass

    def join(self, timeout=None):
        self._stop.set()
        super(BitsSessionCleaner, self).join(timeout)

class BitsServer(BaseHTTPServer.HTTPServer):

    def __init__(self, config, *args, **kw):
        BaseHTTPServer.HTTPServer.__init__(self, *args, **kw)

        self.maxFileSize = 100000000                                                # Default maximum file sise
        self.logfile = "example.log"                                                # Default log file location

        if "maxfilesize" in config:
            self.maxFileSize = config["maxfilesize"]

        if "logfile" in config:
            self.logfile = config["logfile"]

        self.paths = config["paths"]
        self.timeout = config["timeout"]
        self.transit = config["transit"]                                            # Transit directory holds temporary sessions whilst uploading

        logging.basicConfig(filename=self.logfile, level=logging.DEBUG,format='%(asctime)s %(message)s')



class BitsHandler(BaseHTTPServer.BaseHTTPRequestHandler):

    bitsUpload1_5guid = "{7df0354d-249b-430f-820d-3d2a9bef4931}"                    # BITS Upload Protocol 1.5, currently the only version around

    def version_string(self):                                                        
        """Return the server software version string."""

        return self.server_version

    def BITSLog(self, level, command, sessionId, message, *args, **kwargs):
        """Just to do a bit of logging - should be replaced with better logging"""

        logMessage = "%s - %s - %s - %s" % (self.client_address[0], command, sessionId, message)
        logging.log(level, logMessage, *args, **kwargs)

    def createBITSSession(self, urlpath, filename):
        """ Create a new session GUID and start a file for this new transfer """

        guid = "{" + str(uuid.uuid1(random.randrange(0, 0xFFFFFFFFFFFF))) + "}"

        try:
            os.makedirs(self.server.transit + "/" + str(guid))                        # Attempt to make the temporary directory for this
            self.BITSLog(logging.INFO, "Create-Session", str(guid), "Creating transit directory for new job with GUID %s" , str(guid))
        except OSError:    
            self.BITSLog(logging.INFO, "Create-Session", str(guid), "Failed to create transit directory for new job with GUID %s", str(guid))        
            raise BitsSessionCreationError("Failed to create transit directory")

        try:
            file = open(self.server.transit + "/" + str(guid) + "/bits-filename", "w")
            file.write(urlpath+"\n")
            file.write(filename+"\n")
            file.close()

            file = open(self.server.transit + "/" + str(guid) + "/bits-data", "w")
            file.close()
        except IOError:
            self.removeSession(guid)                                                # Remove any part of the session it it exists
            raise BitsSessionCreationError("Error in writing to transit directory")    # Error in setting up session - return None

        return guid    

    def checkAuthorisedUpload(self):
        """ This is a stub method. It may be worth checking other criteria as well """

        # For now, just return true
        return True

    def sendBITSError(self, httpCode, bitsError, bitsErrorContext):
        """ Sends an HTTP error message but add the extra BITS error messages and context """

        self.protocol_version = "HTTP/1.1"
        self.server_version = "PyBUP"

        self.send_response(httpCode)
        self.send_header("Pragma", "no-cache")
        self.send_header("BITS-packet-type", "Ack")
        self.send_header("BITS-Error", bitsError)
        self.send_header("BITS-Error-Context", bitsErrorContext)
        self.end_headers()

    def existsSession(self, sessionId):
        """ Tests for the existence of a session in the transit folder, returning True if present False if not """

        if os.path.isdir(self.server.transit + "/" + sessionId):
            return True
        return False

    def saveFragment(self, sessionId, start, stop, data):
        """ Attempts to save the data fragment to the sessionId from start to stop. Copes with overlaps in ranges """
        """ However, ranges must be contiguous or overlapping, no gaps, and gaps lead to throwing a BitsFragmentSyncError """

        dataFile = self.server.transit + "/" + sessionId + "/bits-data"

        nextExpectedByte = os.path.getsize(dataFile)

        if (start > nextExpectedByte):                                                # Out-of-sync, must send a 416
            raise BitsFragmentSyncError(nextExpectedByte)                            # Raise error with the next expected byte

        try:
            if start == nextExpectedByte:
                f = open(dataFile, 'ab')
                f.write(data)
                f.close
            elif start < nextExpectedByte and stop >= nextExpectedByte:                # There is an overlap, save just the bits needed
                startOffset = nextExpectedByte - start
                f = open(dataFile, 'ab')
                f.write(data[startOffset:])
                f.close
        except IOError:
            raise BitsFragmentSyncError(nextExpectedByte)                            # Failed to write to the file for some reason....

    def processFragment(self):
        """Respond to a Fragment message. Request defined in MC-BUP 2.2.6 """

        if not "Bits-Session-Id" in self.headers:
            self.BITSLog(logging.ERROR, "Fragment", "*", "No BITS Session ID received")
            self.sendBITSError(500, "0x8020001F", "0x5")
        elif not self.existsSession(self.headers["Bits-Session-Id"]):
            self.BITSLog(logging.ERROR, "Fragment", "*", "Specified BITS Session not known on server")
            self.sendBITSError(500, "0x8020001F", "0x5")
        elif not "Content-Range" in self.headers:                                    # Content-Range must be present [ MC-BUP 2.2.6.1 ]
            self.BITSLog(logging.ERROR, "Fragment", "*", "No Content-Range header received")
            self.sendBITSError(400, "0x80070057", "0x5")
        else:
            contentLength = int(self.headers["Content-Length"])
            sessionGUID = self.headers["Bits-Session-Id"]

            contentRange = self.headers["Content-Range"]

            rangeComponents = re.match("bytes (\d+)-(\d+)\/(\*|\d+)", contentRange)
            if not rangeComponents:
                self.BITSLog(logging.ERROR, "Fragment", sessionGUID, "Incorrect Content-Range format: %s", contentRange)

                self.sendBITSError(400, "0x80070057", "0x5")                        # Range not specified in the expected format
            else:
                startByte = int(rangeComponents.group(1))
                stopByte = int(rangeComponents.group(2))
                totalFileLength = rangeComponents.group(3)

                # Testing showed that despite what MC-BUP says, BITS Server on 2008R2 did NOT have a 13Mb max - tested uploading
                # a 128Mb fragment which uploaded on the BITS Server with 
                # So.... not going to put in a max fragment size here (HTTP error 413) unless I can see something else implementing it

                if int(totalFileLength) > self.server.maxFileSize:                    # If file size is too large, tell the client to go away
                    self.BITSLog(logging.ERROR, "Fragment", sessionGUID, "File size greater than maximum permitted")
                    self.sendBITSError(500, "0x80200020", "0x5")

                elif (startByte > stopByte or                                         # Start is greater than stop
                    (stopByte - startByte + 1) > contentLength or                     # Fragment range is greater than datalen            
                    (totalFileLength != "*" and int(totalFileLength) < stopByte + 1)):        # Stop byte is outside range of total file

                    self.BITSLog(logging.ERROR, "Fragment", sessionGUID, "Content-Range invalid: %d-%d/%s %d", startByte, stopByte, totalFileLength, contentLength)
                    self.sendBITSError(400, "0x80070057", "0x5")
                else:

                    try:
                        self.saveFragment(sessionGUID, startByte, stopByte, self.rfile.read(contentLength))
                        self.BITSLog(logging.INFO, "Fragment", sessionGUID, "Saving bytes %d %d", startByte, stopByte)

                        self.protocol_version = "HTTP/1.1"
                        self.server_version = "PyBUP"

                        self.send_response(200)
                        self.send_header("Pragma", "no-cache")
                        self.send_header("BITS-packet-type", "Ack")                    # Mandatory for all BITS responses
                        self.send_header("Content-Length", "0")                        # Must be 0
                        self.send_header("Bits-Session-Id", sessionGUID)            # GUID for client to use in upload
                        self.send_header("BITS-Received-Content-Range", str(stopByte+1))    # Must tell the client to send last + 1 [ MC-BUP 2.2.7.2 ]
                        self.end_headers()
                    except BitsFragmentSyncError as e:
                        self.BITSLog(logging.WARNING, "Fragment", sessionGUID, "Out of sync, asking for byte: %d", str(e.nextExpectedByte))

                        self.protocol_version = "HTTP/1.1"
                        self.server_version = "PyBUP"

                        self.send_response(416)
                        self.send_header("Pragma", "no-cache")
                        self.send_header("BITS-packet-type", "Ack")                    # Mandatory for all BITS responses
                        self.send_header("Content-Length", "0")                        # Must be 0
                        self.send_header("Bits-Session-Id", sessionGUID)            # GUID for client to use in upload
                        self.send_header("BITS-Received-Content-Range", str(e.nextExpectedByte))    # Must tell the client to send last + 1 [ MC-BUP 2.2.7.2 ]
                        self.end_headers()

    def isValidUploadDirectory(self, dirPath):
        """ Looks in the config of the server for whether the specified upload directory is valid and returns True if this is the case  """

        if dirPath in self.server.paths:
            return True
        return False

    def isExistingDirectory(self, dirPath, uploadFile):
        """ Checks to see whether the specified file and path actually represent an existing directory on the server """
        """ Returns True if this is the case, False otherwise """

        if dirPath in self.server.paths:
            rootPath = self.server.paths[dirPath]
            serverFileLocation = rootPath + "/" + uploadFile                        # Form the path to the file on the server

            if os.path.isdir(serverFileLocation):
                return True

        return False

    def isValidFilename(self, filename):
        """ Checks to see whether the upload filename in the URI contains only valid characters. If so returns True else False """

        if re.search('[\\/:\*\?"<>|]', filename):                                    # Does the filename contain invalid characters
            return False
        return True

    def processCreateSession(self):
        """Respond to a Create-Session message. Request defined in MC-BUP 2.2.2 """    

        if "BITS-Supported-Protocols" in self.headers:                                # Create-Session messages must specify BITS-Supported-Protocols
            supportedProtocols = self.headers["BITS-Supported-Protocols"]            # Confirm that this protocol is an acceptable protocol version
            supportedProtocolsList = supportedProtocols.split(" ")    

            (dirPath, uploadFile) = self.path.rsplit('/', 1)                        # Split into the dir and file

            if not self.isValidFilename(uploadFile):                                # Does this filename contain invalid characters?
                self.BITSLog(logging.ERROR, "Create-Session", "*", "Specified filename is illegal")
                self.sendBITSError(400, "0x80070057", "0x5")
            elif not self.isValidUploadDirectory(dirPath):
                self.BITSLog(logging.ERROR, "Create-Session", "*", "Specified upload directory is not enabled for upload")
                self.sendBITSError(501, "0x80070005", "0x5")                        # Uploads not enabled for this virtual directory
            elif self.isExistingDirectory(dirPath, uploadFile):
                self.BITSLog(logging.ERROR, "Create-Session", "*", "Specified file for upload is an existing directory")
                self.sendBITSError(400, "0x80070057", "0x5")
            elif not (self.bitsUpload1_5guid in supportedProtocolsList):
                self.BITSLog(logging.ERROR, "Create-Session", "*", "BITS protocol version is not supported")
                self.sendBITSError(400, "0x80070057", "0x5")                        # Unsupported protocol version number, send a 400 error
            elif not self.checkAuthorisedUpload():
                self.BITSLog(logging.ERROR, "Create-Session", "*", "Upload is not authorised")
                self.sendBITSError(403, "0x80070005", "0x5")                        # Upload attempt to unauthorised location or cannot be overwritten
            else:
                # To do: pass filename to save it
                try:
                    sessionGUID = self.createBITSSession(dirPath, uploadFile)        # Attempt to create a new session, error thrown if fails

                    self.protocol_version = "HTTP/1.1"
                    self.server_version = "PyBUP"

                    self.send_response(200)
                    self.send_header("Pragma", "no-cache")
                    self.send_header("BITS-packet-type", "Ack")                        # Mandatory for all BITS responses
                    self.send_header("Accept-Encoding", "Identity")                    # The only acceptable encoding for BITS Upload
                    self.send_header("Content-Length", "0")                            # Must be 0
                    self.send_header("Bits-Protocol", self.bitsUpload1_5guid)        # BITS Upload Protocol 1.5
                    self.send_header("Bits-Session-Id", sessionGUID)                # GUID for client to use in upload
                    self.end_headers()

                    self.BITSLog(logging.INFO, "Create-Session", sessionGUID, "New session for %s", self.path)

                except BitsSessionCreationError as e:                                # Catch error in producing a new session with GUID
                    self.sendBITSError(403, "0x80070005", "0x5")                    # Return error if upload cannot be created in transit area
        else:
            self.BITSLog(logging.ERROR, "Create-Session", "*", "No protocol version specified")
            self.sendBITSError(400, "0x80070057", "0x5")

    def processPing(self):
        """ Simple response to a BITS Ping """
        self.protocol_version = "HTTP/1.1"
        self.server_version = "PyBUP"

        self.send_response(200)
        self.send_header("Pragma", "no-cache")
        self.send_header("BITS-packet-type", "Ack")                                    # Mandatory for all BITS responses
        self.send_header("Content-Length", "0")                                        # Must be 0
        self.end_headers()

        self.BITSLog(logging.INFO, "Ping", "*", "Ping Response")

    def moveSessionToDestination(self, sessionId):
        """ Called once the Close-Session message has been received, to move it to the final location """

        sessionPath = self.server.transit + "/" + sessionId

        dirpath = ""
        filename = ""

        with open(sessionPath + "/bits-filename") as tempFile:
            dirpath = tempFile.readline().rstrip('\n')
            filename = tempFile.readline().rstrip('\n')

        if not dirpath in self.server.paths:                                        # Ensure this is moving to a genuine virtual directory
            self.BITSLog(logging.ERROR, "Close-Session", sessionId, "Invalid virtual directory: %s", dirpath)
            self.sendBITSError(403, "0x80070005", "0x5")
            return False
        elif not self.isValidFilename(filename):                                    # Does this filename contain invalid characters?
            self.BITSLog(logging.ERROR, "Close-Session", sessionId, "Specified filename is illegal: %s", filename)
            self.sendBITSError(400, "0x80070057", "0x5")
            return False
        else:
            try:
                shutil.move(sessionPath + "/bits-data", self.server.paths[dirpath] + "/" + filename)
            except shutil.Error:
                self.BITSLog(logging.ERROR, "Close-Session", sessionId, "Could not move completed session to final directory")
                self.sendBITSError(403, "0x80070005", "0x5")
                return False
        return True

    def removeSession(self, guid):
        """ Remove the temporary session folder from the transit directory """

        path = self.server.transit + "/" + guid

        self.BITSLog(logging.INFO, "Remove-Session", guid, "Removing session directory %s:", path)

        try:        
            shutil.rmtree(path)                                                        # Attempt to remove the entire tree of the guid
        except OSError:
            pass

    def processCloseSession(self):
        """ Process the Close-Session message, received from the client when an upload is complete """

        if not "Bits-Session-Id" in self.headers:                                    # Check to see that a session ID is specified
            self.BITSLog(logging.ERROR, "Close-Session", "*", "No BITS Session ID received")
            self.sendBITSError(500, "0x8020001F", "0x5")
        elif not self.existsSession(self.headers["Bits-Session-Id"]):                # And that it is an existing session
            self.BITSLog(logging.ERROR, "Close-Session", "*", "Specified BITS Session not known on server")
            self.sendBITSError(500, "0x8020001F", "0x5")
        else:
            sessionGuid = self.headers["Bits-Session-Id"]

            if self.moveSessionToDestination(sessionGuid):                            # Move file to final directory            

                self.BITSLog(logging.INFO, "Close-Session", sessionGuid, "Closing session %s", sessionGuid)
                self.removeSession(sessionGuid)                                        # Clean up temporary files

                self.protocol_version = "HTTP/1.1"
                self.server_version = "PyBUP"            
                self.send_response(200)
                self.send_header("Pragma", "no-cache")
                self.send_header("BITS-packet-type", "Ack")                            # Mandatory for all BITS responses
                self.send_header("BITS-Session-Id", sessionGuid)
                self.send_header("Content-Length", "0")                                # Must be 0
                self.end_headers()


    def processCancelSession(self):
        """ Received when the client wishes to cancel an upload """

        if not "Bits-Session-Id" in self.headers:                                    # Check to see that a session ID is specified
            self.BITSLog(logging.ERROR, "Cancel-Session", "*", "No BITS Session ID received")
            self.sendBITSError(500, "0x8020001F", "0x5")
        elif not self.existsSession(self.headers["Bits-Session-Id"]):                # And that it is an existing session
            self.BITSLog(logging.ERROR, "Cancel-Session", "*", "Specified BITS Session not known on server")
            self.sendBITSError(500, "0x8020001F", "0x5")
        else:            
            sessionGuid = self.headers["Bits-Session-Id"]

            self.BITSLog(logging.INFO, "Cancel-Session", sessionGuid, "Cancelling session: %s", sessionGuid)
            self.removeSession(sessionGuid)

            self.protocol_version = "HTTP/1.1"
            self.server_version = "PyBUP"            
            self.send_response(200)
            self.send_header("Pragma", "no-cache")
            self.send_header("BITS-packet-type", "Ack")                                # Mandatory for all BITS responses
            self.send_header("BITS-Session-Id", sessionGuid)
            self.send_header("Content-Length", "0")                                    # Must be 0
            self.end_headers()


    def do_BITS_POST(self):
        """Respond to a BITS Upload reuest message"""

        # First thing is to determine what type of BITS reuest this is. This is specified in the BITS-Packet-Type header
        # Ensure that this, the content length and the HTTP version is 1.1 as these are mandatory [ MC-BUP 2.2.1 - 2.2.1.2 ]

        if ("BITS-Packet-Type" in self.headers and 
            self.request_version == "HTTP/1.1" and 
            "Content-Length" in self.headers):

            packetType = self.headers["BITS-Packet-Type"]                            # Switch based on packetType

            if packetType == "Create-Session":                                        #  [ MC-BUP 2.2.2 ]
                self.processCreateSession()
            elif packetType == "Fragment":                                            #  [ MC-BUP 2.2.6 ]
                self.processFragment()
            elif packetType == "Ping":
                self.processPing()
            elif packetType == "Close-Session":                                        #  [ MC-BUP 2.2.8 ]
                self.processCloseSession()
            elif packetType == "Cancel-Session":
                self.processCancelSession()
            else:
                self.BITSLog(logging.ERROR, "*", "*", "Unknown packet-type: %s", packetType)
                self.sendBITSError(400, "0x80070057", "0x5")                        # If an unknown packetType, correct response is a 400
        elif not "Content-Length" in self.headers:
            self.BITSLog(logging.ERROR, "*", "*", "No Content-length")
            self.sendBITSError(411, "0x80070057", "0x5")                            # See [ MC-BUP 3.2.5.2.3]
        else:
            self.BITSLog(logging.ERROR, "*", "*", "No Packet-Type header")
            self.sendBITSError(500, "0x80070585", "0x5")                            # Failure to set these headers correctly results in a 500 error
                                                                                    # This isn't as specifed in MC-BUP but appears to be what happens    
        logger = logging.getLogger()
        logger.handlers[0].flush()

    # def do_HEAD(self):
    # Implement as per standard web server

    # def do_GET(self):
    # Implement as per standard web server

if __name__ == '__main__':

    # For the purposes of this demo, the config is just stuck in a dirty dictionary

    bitsConfig = { "transit" : "./transit",
                    "paths" : {"/upload/bits" : ".",
                                "/post/files" : "."},
                    "timeout" : 604800,                                                # Set the session timeout to 7 days
                    "logfile" : "bits.log",
                    "port" : 81 }

    httpd = BitsServer(bitsConfig, ("", bitsConfig["port"]), BitsHandler)

    cleaner = BitsSessionCleaner(bitsConfig)
    cleaner.start()

    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        pass

    cleaner.join()
    httpd.server_close()